#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
计算[成本路径](https://pro.arcgis.com/zh-cn/pro-app/tool-reference/spatial-analyst/how-the-cost-distance-tools-work.htm)的算法
- 目前仅支持矢量点与成本栅格的成本路径计算
- 计算时间较长,考虑进一步优化算法
"""

import sys
import math
import numpy as np
from osgeo import gdal, ogr, osr

def get_neighbors(coordinate, max_range):
    """计算某一坐标点的相邻坐标
        默认该点周围的 8 个点,同时也会考虑范围的影响
        coordinate 给定点的坐标
        max_range 数组最大范围
    """
    max_x = max_range[0]
    max_y = max_range[1]

    result = []
    for x in range(coordinate[0] - 1, coordinate[0] + 2):
        for y in range(coordinate[1] - 1, coordinate[1] + 2):
            if x >= 0 and x < max_x and y >= 0 and y < max_y:
                if x == coordinate[0] and y == coordinate[1]:
                    continue

                result.append((x, y))

    return result

def basic_cost(cost_array, start_point, end_point):
    """计算相邻结点成本
        cost_array 成本数组
        (start_x, start_y) 起点序号
        (end_x, end_y) 终点序号
    """
    start_x, start_y = start_point[0], start_point[1]
    end_x, end_y = end_point[0], end_point[1]

    p1 = np.array([start_x, start_y])
    p2 = np.array([end_x, end_y])

    distance = np.sum(np.abs(p2 - p1))
    # pprint(distance)

    assert distance > 0 and distance <= 2

    if distance == 1:
        cost = (cost_array[start_x, start_y] + cost_array[end_x, end_y]) / 2
    elif distance == 2:
        cost = math.sqrt(2) * (cost_array[start_x, start_y] + cost_array[end_x, end_y]) / 2

    return cost

def accumulative_cost(cost_array, input_raster, start_point):
    """计算与给定点相邻的其他所有点的成本距离
        cost_array 成本栅格
        input_raster 输入栅格
        start_point 给定计算的起始点
    """
    source_cost = input_raster[start_point[0], start_point[1]]
    shape = input_raster.shape

    for end_point in get_neighbors(start_point, shape):
        inactive_cost = input_raster[end_point[0], end_point[1]]
        active_cost = basic_cost(cost_array, start_point, end_point)

        if np.isnan(inactive_cost) or inactive_cost  > active_cost + source_cost:
            input_raster[end_point[0], end_point[1]] = active_cost + source_cost

def calculate_cost_distance(cost_array, input_raster):
    """遍历整个栅格计算距离成本
        cost_array 成本栅格
        input_raster 输入栅格
    """
    accumulative_minimal_array = input_raster.copy()

    while not np.isnan(np.nanmin(accumulative_minimal_array)):
        input_nanmin = np.nanmin(accumulative_minimal_array)
        nanmin_index = np.where(input_raster == input_nanmin)

        for coordinate in zip(nanmin_index[0], nanmin_index[1]):
            accumulative_cost(cost_array, input_raster, coordinate)

        accumulative_minimal_array = input_raster.copy()
        accumulative_minimal_array[accumulative_minimal_array <= input_nanmin] = np.nan

def coord2pixelOffset(raster, coordinates):
    """根据给定的空间坐标,计算该点相对栅格的像素偏移
        raster (gdal.Dataset) 成本栅格 
        coordinates 空间坐标
    """
    geotransform = raster.GetGeoTransform()
    originX = geotransform[0]
    originY = geotransform[3]
    pixelWidth = geotransform[1]
    pixelHeight = geotransform[5]

    results = []
    for x, y in coordinates:
        xOffset = int((x - originX)/pixelWidth)
        yOffset = int((y - originY)/pixelHeight)
        results.append((xOffset, yOffset))
    return results

def main(shapefile, rasterfile, out_tif='OUT.tif'):
    """读取文件,处理坐标和数据,然后计算成本路径并保存
        shapefile 输入矢量路径
        rasterfile 输入成本栅格路径
        out_tif 输出文件名
    """
    raster = gdal.Open(rasterfile)
    band = raster.GetRasterBand(1)

    shape = ogr.Open(shapefile)
    layer = shape.GetLayer()

    # 获取栅格空间信息
    transform = raster.GetGeoTransform()
    x_origin = transform[0]
    y_origin = transform[3]
    pixel_width = transform[1]
    pixel_height = transform[5]

    # 将矢量数据空间坐标系投影到栅格
    source_sr = layer.GetSpatialRef()
    target_sr = osr.SpatialReference()
    target_sr.ImportFromWkt(raster.GetProjectionRef())
    coord_trans = osr.CoordinateTransformation(source_sr, target_sr)

    points_x = []
    points_y = []

    for f_id in range(layer.GetFeatureCount()):
        feature = layer.GetFeature(f_id)

        geom = feature.GetGeometryRef()
        geom.Transform(coord_trans)

        if geom.GetGeometryName() == 'POINT':
            points = geom.GetPointCount()

            for p in range(points):
                lon, lat, z = geom.GetPoint(p)
                points_x.append(lon)
                points_y.append(lat)

        else:
            raise ValueError('unsupport geometry type')

    # 注意:留出两个像素点的位置作为缓冲区
    x_min = min(points_x) - 2 * pixel_width
    x_max = max(points_x) + 2 * pixel_width
    y_min = min(points_y) - 2 * pixel_width
    y_max = max(points_y) + 2 * pixel_width

    # 指定需要在栅格上读出的偏移及行列像素位 
    x_offset = int((x_min - x_origin) / pixel_width)
    y_offset = int((y_origin - y_max) / pixel_width)
    x_count = int((x_max - x_min) / pixel_width)
    y_count = int((y_max - y_min) / pixel_width)

    # 使用内存栅格来保存矢量点外接矩形与成本栅格相切后的数据
    target_ds = gdal.GetDriverByName('MEM').Create('', x_count, y_count, 1, gdal.GDT_Float32)
    target_ds.SetGeoTransform((x_min, pixel_width, 0, y_max, 0, pixel_height))
    target_band = target_ds.GetRasterBand(1)

    raster_array = band.ReadAsArray(x_offset, y_offset, x_count, y_count)

    target_band.WriteArray(raster_array)

    # 使对象栅格的投影坐标与原始栅格保持一致
    raster_srs = osr.SpatialReference()
    raster_srs.ImportFromWkt(raster.GetProjectionRef())
    target_ds.SetProjection(raster_srs.ExportToWkt())

    # target_band.FlushCache()

    # 将与矢量点对应的输入栅格坐标标记为 0, 其余为 np.nan
    shp_points = coord2pixelOffset(target_ds, list(zip(points_x, points_y)))
    target_array = target_band.ReadAsArray().astype(np.float)

    # pprint(target_array.shape)
    input_array = np.empty(target_array.shape)
    input_array[:]  = np.nan

    # pprint(shp_points)

    for col, row in shp_points:
        input_array[row, col] = 0

    # 计算成本路径
    calculate_cost_distance(target_array, input_array)
    # pyplot.imshow(input_array)

    # 将计算结构写入新的栅格
    output_ds = gdal.GetDriverByName('GTiff').Create(out_tif, x_count, y_count, 1, gdal.GDT_Float32)
    output_ds.SetGeoTransform((x_min, pixel_width, 0, y_max, 0, pixel_height))
    output_band = output_ds.GetRasterBand(1)

    output_band.WriteArray(input_array)
    output_ds.SetProjection(raster_srs.ExportToWkt())
    output_band.FlushCache()


